import torch
import torch.nn as nn
import config


# 基本模型，包含了训练和测试方法
class BaseModel(nn.Module):
    def __init__(self, is_loaded=False, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def load(self):
        self.load_state_dict(torch.load(config.save_name, map_location=config.device))

    def train_step(self, losser, batch):
        images, boxes, labels = batch
        inputs, boxes, labels = images.to(config.device), boxes.to(config.device), labels.to(
            config.device)  # 放在GPU上加速训练
        return losser.train_step(images, boxes, labels, 1)

    # 测试模型正确率
    def evaluate(self, data):
        val_loss = .0
        with torch.no_grad():
            for batch in data:
                images, boxes, labels = batch
                inputs, boxes, labels = images.to(config.device), boxes.to(config.device), labels.to(
                    config.device)  # 放在GPU上加速训练
                _, _, _, _, val_total = self.forward(images, boxes, labels, 1)
                val_loss += val_total.item()

        return val_loss
